
import requests
# from .base import GroupItem


class ItemOutput:

    def __init__(self,imagepath=None,text=None):
        self.imagepath = imagepath
        self.text = text

    def open_to_json(self):
        if self.imagepath:
            return {
                "image": open(self.imagepath,"r"),
                "text": self.text
            }
        else:
            return { "text": self.text }
 

class GroupItem:

    IMAGETAG = "<ImageHere>"

    def __init__(self,image=None,**kwargs):
        self.kwargs = kwargs
        self.imagepath = None
        if image:
            self.imagepath = image
            self.kwargs["image"] = self.IMAGETAG
        self.keys = list(kwargs.keys())
    
    def get_image_path(self):
        return self.imagepath
    
    def __extract_kwargs(self,template:str):
        kwargs = {}
        for key in self.keys:
            item = "{" + str(key) + "}"
            if item in template:
                kwargs[key] = self.kwargs[key]
        return kwargs
    
    def __make(self,template:str,allmatch=True):
        kwargs = self.__extract_kwargs(template)
        if allmatch:
            return template.format_map(kwargs)
        else:
            result = str(template)
            for key,value in kwargs.items():
                item = "{" + str(key) + "}"
                result = result.replace(item,str(value))
            return result

    def Make(self,template:str,allmatch=True):
        return self.__make(template,allmatch)
    
    def OutJson(self,template:str,allmatch=True):
        text = self.__make(template,allmatch)
        return ItemOutput(self.imagepath,text)
    


# if __name__ == "__main__":

#     template = "{image}, {name} is {age}"

#     item = GroupItem(name='sdf',age='1234',image="main.py")
#     a = item.Make(template)
#     print(a)

#     a = item.OutJson(template)
#     print(a.open_to_json())




class LLMClient:

    def __init__(self,url:str="127.0.0.1:7002") -> None:
        self.url = f"http://{url}/predict"
        pass
    
    def Predict(self, text:str, imagepath:str=None):
        item = GroupItem(image=imagepath,query=text)

        if imagepath:
            template = "###Human: <Img>{image}</Img> {query} ###Assistant:"
            text = item.Make(template)
            data = {
                "text" : (None, text),
                "json" : (None, "")
            }
            files_list = [
                ("images", open(imagepath, "rb"))
            ]
            # print(files_list)
            response = requests.post(self.url, data=data, files=files_list)
        else:
            template = "###Human: {query} ###Assistant:"
            text = item.Make(template)
            data = {
                "text" : (None, text),
                "json" : (None, "")
            }
            response = requests.post(self.url, data=data)

        json_output = response.json()
        # print(json_output)
        output = str(json_output["result"])
        # 打印返回的结果
        # print(response.json())
        print(output)

        return output



if __name__ == "__main__":
    # client = LLMClient("127.0.0.1:5000")
    client = LLMClient("81.70.84.40:64091")
    # client.Predict("what is the result of 1+1?")
    client.Predict("what is in this image?","teaser.png")



